package com.example.aa175.fcn;

import android.Manifest;
import android.app.Activity;
import android.content.Context;
import android.content.Intent;
import android.content.pm.PackageManager;
import android.graphics.Bitmap;
import android.graphics.BitmapFactory;
import android.graphics.Color;
import android.graphics.drawable.BitmapDrawable;
import android.net.Uri;
import android.os.Bundle;
import android.os.Environment;
import androidx.annotation.NonNull;
import androidx.annotation.Nullable;
import androidx.core.app.ActivityCompat;
import androidx.core.content.ContextCompat;
import androidx.appcompat.app.AppCompatActivity;
import android.util.Log;
import android.view.View;
import android.widget.Button;
import android.widget.ImageView;
import android.widget.Toast;

import org.pytorch.IValue;
import org.pytorch.Module;
import org.pytorch.Tensor;
import org.pytorch.torchvision.TensorImageUtils;

import java.io.BufferedInputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.List;

// 主活动类
public class MainActivity extends AppCompatActivity {
    // 创建pytorch模型
    private Module segModule;

    // 定义ImageView对象show_image展示的图片的路径
    private String showImagePath;
    // 定义相机拍摄图像的路径
    private String cameraImagePath;
    // 定义从相册选择图像的路径
    private String albumImagePath;

    private static final int START_CAMERA_CODE = 1111;   // 开启相机的请求码
    private static final int START_ALBUM_CODE = 1112;    // 开启相册的请求码
    private static final int REQUIRE_PERMISSION_CODE = 111; // 权限请求码

    // 创建用于开启相机的按钮
    private Button btnStartCamera;
    // 创建用于调用系统相册的按钮
    private Button btnStartAlbum;
    // 创建用于保存ImageView内容的按钮
    private Button btnSaveImage;
    // 创建用于启动pytorch前向推理的按钮
    private Button btnStartPredict;

    // 创建ImageView对象
    private ImageView showOriginalImageView;  // 展示原图
    private ImageView showMaskImageView;      // 展示分割掩码
    private ImageView showForegroundImageView; // 展示分割后的前景图像

    @Override
    protected void onCreate(Bundle savedInstanceState) { // 创建活动
        super.onCreate(savedInstanceState);  // 调用父类的onCreate方法
        requestPermissions();  // 请求权限
        setContentView(R.layout.activity_main);  // 设置布局文件

        // 绑定开启相机的按钮
        btnStartCamera = (Button) findViewById(R.id.button_start_camera);// 获取布局文件中的button_start_camera
        btnStartCamera.setOnClickListener(new StartCameraOnClickListener());//创建StartCameraOnClickListener对象，然后将其传入setOnClickListener方法中
        // 绑定调用相册的按钮
        btnStartAlbum = (Button) findViewById(R.id.button_start_album);// 获取布局文件中的button_start_album
        btnStartAlbum.setOnClickListener(new StartAlbumOnClickListener());//创建StartAlbumOnClickListener对象，然后将其传入setOnClickListener方法中
        // 绑定保存图片的按钮
        btnSaveImage = (Button) findViewById(R.id.button_save_image);// 获取布局文件中的button_save_image
        btnSaveImage.setOnClickListener(new SaveImageOnClickListener());//创建SaveImageOnClickListener对象，然后将其传入setOnClickListener方法中
        // 绑定启动pytorch前向推理的按钮
        btnStartPredict = (Button) findViewById(R.id.button_start_predict);// 获取布局文件中的button_start_predict
        btnStartPredict.setOnClickListener(new StartPredictOnClickListener());//创建StartPredictOnClickListener对象，然后将其传入setOnClickListener方法中

        // 绑定我们在activity_main.xml中定义的Image_View
        showOriginalImageView = (ImageView) findViewById(R.id.image_view_show_original);// 获取布局文件中的image_view_show_original
        showMaskImageView = (ImageView)findViewById(R.id.image_view_show_predict_mask);// 获取布局文件中的image_view_show_predict_mask
        showForegroundImageView = (ImageView)findViewById(R.id.image_view_show_body);// 获取布局文件中的image_view_show_body
    }

    class StartCameraOnClickListener implements View.OnClickListener {  // 创建StartCameraOnClickListener类，实现View.OnClickListener接口
        @Override
        public void onClick(View v) {
            // 开启相机，返回相机拍摄图像的路径，传入的请求码是START_CAMERA_CODE
            cameraImagePath = Utils.startCamera(MainActivity.this, START_CAMERA_CODE);
            // 图片保存成功后，更新相册
            if (cameraImagePath != null && !cameraImagePath.isEmpty()) {
                // 图片保存成功，立即通知相册更新
                Utils.updateGallery(MainActivity.this, cameraImagePath);
            }
        }
    }

    class StartAlbumOnClickListener implements View.OnClickListener {  // 创建StartAlbumOnClickListener类，实现View.OnClickListener接口
        @Override
        public void onClick(View v) {
            // 调用相册，传入的请求码是START_ALBUM_CODE
            Utils.startAlbum(MainActivity.this, START_ALBUM_CODE);
        }
    }

    class SaveImageOnClickListener implements View.OnClickListener {  // 创建SaveImageOnClickListener类，实现View.OnClickListener接口
        @Override
        public void onClick(View v) {
            // 获取show_image展示的图片
            Bitmap bitmap = ((BitmapDrawable) showForegroundImageView.getDrawable()).getBitmap();
            // 通过saveImage函数生成图像的保存路径
            String imageSavePath = saveImage(bitmap, 100);
            // 提示用户图片已经被保存
            if (imageSavePath != null) {
                Toast.makeText(MainActivity.this, "save to: " + imageSavePath, Toast.LENGTH_LONG).show();
            }
            // 图片保存成功后，更新相册
            if (imageSavePath != null && !imageSavePath.isEmpty()) {
                // 图片保存成功，立即通知相册更新
                Utils.updateGallery(MainActivity.this, imageSavePath);
            }
        }
    }

    class StartPredictOnClickListener implements View.OnClickListener {  // 创建StartPredictOnClickListener类，实现View.OnClickListener接口
        @Override
        public void onClick(View v) {
            segmentationForMask(showImagePath);
        } // 调用segmentationForMask函数，传入的参数是showImagePath
    }

    private static String saveImage(Bitmap bmp, int quality) { // 创建saveImage函数，用于保存图像
        if (bmp == null) {
            return null;
        }  // 如果bmp为空，则返回null

        // 获取外部存储的公共图片目录
        File appDir = Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_PICTURES);

        if (appDir == null) {
            return null;
        }  // 如果appDir为空，则返回null

        String fileName = System.currentTimeMillis() + ".jpg";  // 生成一个唯一的文件名
        File file = new File(appDir, fileName);  // 创建一个文件，传入的参数是appDir和fileName
        FileOutputStream fos = null;  // 创建一个FileOutputStream对象

        try {
            fos = new FileOutputStream(file);  // 创建一个FileOutputStream对象，传入的参数是file
            bmp.compress(Bitmap.CompressFormat.JPEG, quality, fos);  // 将图像压缩为JPEG格式，压缩质量为quality，然后写入到fos中
            fos.flush();  // 将缓冲区中的数据全部写出到输出流中
            return file.getAbsolutePath();  // 返回图像的保存路径

        } catch (FileNotFoundException e) {  // 捕获FileNotFoundException异常
            e.printStackTrace();  // 打印异常信息
        } catch (IOException e) {  // 捕获IOException异常
            e.printStackTrace();  // 打印异常信息
        } finally {
            if (fos != null) {
                try {
                    fos.close(); // 关闭fos
                } catch (IOException e) { // 捕获IOException异常
                    e.printStackTrace(); // 打印异常信息
                }
            }  // 如果fos不为空，则关闭fos
        }
        return null;

    }

    // 请求本案例需要的三种权限
    private void requestPermissions() {
        // 定义容器，存储我们需要申请的权限
        List<String> permissionList = new ArrayList<>();
        // 检测应用是否具有CAMERA的权限
        if (ContextCompat.checkSelfPermission(this, Manifest.permission.CAMERA) != PackageManager.PERMISSION_GRANTED) {
            permissionList.add(Manifest.permission.CAMERA);
        }
        // 检测应用是否具有READ_EXTERNAL_STORAGE权限
        if (ContextCompat.checkSelfPermission(this, Manifest.permission.READ_EXTERNAL_STORAGE) != PackageManager.PERMISSION_GRANTED) {
            permissionList.add(Manifest.permission.READ_EXTERNAL_STORAGE);
        }

        // 检测应用是否具有WRITE_EXTERNAL_STORAGE权限
        if (ContextCompat.checkSelfPermission(this, Manifest.permission.WRITE_EXTERNAL_STORAGE) != PackageManager.PERMISSION_GRANTED) {
            permissionList.add(Manifest.permission.WRITE_EXTERNAL_STORAGE);
        }

        // 如果permissionList不为空，则说明前面检测的三种权限中至少有一个是应用不具备的
        // 则需要向用户申请使用permissionList中的权限
        if (!permissionList.isEmpty()) {
            ActivityCompat.requestPermissions(this, permissionList.toArray(new String[permissionList.size()]), REQUIRE_PERMISSION_CODE);
        }
    }

    @Override
    public void onRequestPermissionsResult(int requestCode, @NonNull String[] permissions, @NonNull int[] grantResults) {
        super.onRequestPermissionsResult(requestCode, permissions, grantResults); // 调用父类的onRequestPermissionsResult方法
        // 判断请求码
        switch (requestCode) {
            // 如果请求码是我们设定的权限请求代码值，则执行下面代码
            case REQUIRE_PERMISSION_CODE:
                if (grantResults.length > 0) {
                    for (int i = 0; i < grantResults.length; i++) {
                        // 如果请求被拒绝，则弹出下面的Toast
                        if (grantResults[i] == PackageManager.PERMISSION_DENIED) {
                            Toast.makeText(this, permissions[i] + " was denied", Toast.LENGTH_SHORT).show();
                        }
                    }
                }
                break;
        }
    }

    @Override
    protected void onActivityResult(int requestCode, int resultCode, @Nullable Intent data) { // 创建onActivityResult函数，用于处理相机和相册返回的结果
        super.onActivityResult(requestCode, resultCode, data); // 调用父类的onActivityResult方法
        if (resultCode == Activity.RESULT_OK) {  // 如果返回结果是RESULT_OK，则执行下面代码
            switch (requestCode) { // 判断请求码
                case START_ALBUM_CODE:  // 如果请求码是START_ALBUM_CODE，则执行下面代码
                    if (data == null) {
                        Log.w("LOG", "user photo data is null"); // 打印日志
                        return;
                    }

                    Uri albumImageUri = data.getData(); // 获取相册返回的图像资源标识
                    albumImagePath = Utils.getPathFromUri(MainActivity.this, albumImageUri);  // 获取图像的路径
                    showImagePath = albumImagePath;  // 将图像的路径赋值给showImagePath
                    Bitmap bitmap_album = Utils.getScaleBitmapByPath(albumImagePath);  // 通过getScaleBitmapByPath函数获取图像的Bitmap对象
                    showOriginalImageView.setImageBitmap(Utils.getScaleBitmapByBitmap(bitmap_album, 448, 448)); // 将图像展示到ImageView上
                    Toast.makeText(MainActivity.this, "album start", Toast.LENGTH_LONG).show();  // 弹出Toast提示用户相册已经启动
                    break;
                case START_CAMERA_CODE:  // 如果请求码是START_CAMERA_CODE，则执行下面代码
                    showImagePath = cameraImagePath;  // 将图像的路径赋值给showImagePath
                    Bitmap bitmap_camera = Utils.getScaleBitmapByPath(cameraImagePath);  // 通过getScaleBitmapByPath函数获取图像的Bitmap对象
                    Toast.makeText(MainActivity.this, "camera start", Toast.LENGTH_LONG).show();  // 弹出Toast提示用户相机已经启动
                    showOriginalImageView.setImageBitmap(Utils.getScaleBitmapByBitmap(bitmap_camera, 448, 448));  // 将图像展示到ImageView上
                    break;
            }
        }
    }

    /**
     * @param imagePath
     */
    //此方法接收一个图像文件路径作为参数，通过方法内部加载模型和处理图像，进行图像分割和展示结果
    private void segmentationForMask(String imagePath) {
        //图像处理过程
        String ptPath = "20_unet.pt";
        // 创建3通道的输出，虽然输出的mask是单通道的，这里将mask看作是3通道的
        int[] inDims = {256, 256, 3};
        int[] outDims = {256, 256, 3};

        int mask_channels = 1;  // mask_channels是输出的通道数
        Bitmap bmp = null;  // 创建Bitmap对象
        Bitmap scaledBmp = null;  // 创建Bitmap对象
        String filePath = ""; // 创建filePath变量
        try {
            filePath = assetFilePath(this, ptPath);  // 获取ptPath的路径
            segModule = Module.load(filePath);  // 加载ptPath路径下的模型
        } catch (Exception e) {
            Log.d("LOG", "can not load pt " + filePath); // 打印日志
        }
        try {
            BufferedInputStream bis = new BufferedInputStream(new FileInputStream(imagePath));  // 创建BufferedInputStream对象
            bmp = BitmapFactory.decodeStream(bis);  // 通过decodeStream函数将BufferedInputStream对象转换为Bitmap对象
            scaledBmp = Bitmap.createScaledBitmap(bmp, inDims[0], inDims[1], true);  // 将bmp缩放为inDims[0] * inDims[1]大小的图像
            bis.close();
        } catch (Exception e) {
            Log.d("LOG", "can not read human image bitmap");  // 打印日志
        }

        float[] meanRGB = {0.0f, 0.0f, 0.0f}; // 创建meanRGB数组
        float[] stdRGB = {1.0f, 1.0f, 1.0f};  // 创建stdRGB数组
        Tensor humanImageTensor = TensorImageUtils.bitmapToFloat32Tensor(scaledBmp,
                meanRGB, stdRGB); // 将scaledBmp转换为Tensor对象

        try {
            Tensor maskTensor = segModule.forward(IValue.from(humanImageTensor)).toTensor();  // 将humanImageTensor传入模型，获取输出的maskTensor
            float[] maskArray = maskTensor.getDataAsFloatArray();  // 将maskTensor转换为float数组

            int index = 0;
            // mask_image用于存放分割掩码，它是3通道的
            float[][][] maskImage = new float[outDims[0]][outDims[1]][3];  // 创建maskImage数组
            for (int j = 0; j < mask_channels; j++) {
                for (int k = 0; k < outDims[0]; k++) {
                    for (int m = 0; m < outDims[1]; m++) {
                        // 对3个通道进行赋值，遍历的时候j只遍历1个通道值，因为输出的outArr是1通道的
                        if(maskArray[index] > 0.5){ // 如果maskArray[index]大于0.5，则将maskImage[k][m][j]、maskImage[k][m][j + 1]、maskImage[k][m][j + 2]赋值为255.0f
                            maskImage[k][m][j] = 255.0f;
                            maskImage[k][m][j + 1] = 255.0f;
                            maskImage[k][m][j + 2] = 255.0f;
                        }
                        else { // 如果maskArray[index]小于等于0.5，则将maskImage[k][m][j]、maskImage[k][m][j + 1]、maskImage[k][m][j + 2]赋值为0.0f
                            maskImage[k][m][j] = 0.0f;
                            maskImage[k][m][j + 1] = 0.0f;
                            maskImage[k][m][j + 2] = 0.0f;
                        }
                        index++;
                    }
                }
            }

            // 将分割掩码展示到ImageView上
            Bitmap maskBitmap = Utils.getBitmap(maskImage, outDims);
            showMaskImageView.setImageBitmap(maskBitmap);

            // foregroundImage用于展示经分割掩码处理后的前景图像
            float[][][] foregroundImage = new float[inDims[0]][inDims[1]][inDims[2]];
            index = 0;

            for (int k = 0; k < outDims[0]; k++) {
                for (int m = 0; m < outDims[1]; m++) {
                    // 获取原图中的像素值
                    int pixel = scaledBmp.getPixel(m, k);
                    int r = Color.red(pixel);   // 获取r通道值
                    int g = Color.green(pixel); // 获取g通道值
                    int b = Color.blue(pixel);  // 获取b通道值
                    // 对3个通道进行赋值，遍历的时候j只遍历1个通道值
                    if(maskArray[index] > 0.5){  // 如果maskArray[index]大于0.5，则将foregroundImage[k][m][0]、foregroundImage[k][m][1]、foregroundImage[k][m][2]赋值为r、g、b
                        foregroundImage[k][m][0] =  r;
                        foregroundImage[k][m][1] = g;
                        foregroundImage[k][m][2] =b;
                    }
                    else {  // 如果maskArray[index]小于等于0.5，则将foregroundImage[k][m][0]、foregroundImage[k][m][1]、foregroundImage[k][m][2]赋值为0、0、0
                        foregroundImage[k][m][0] =  0;
                        foregroundImage[k][m][1] = 0;
                        foregroundImage[k][m][2] =0;
                    }
                    index++;
                }
            }

            // 将分割后的前景图像展示到ImageView上
            Bitmap foregroundBitmap = Utils.getBitmap(foregroundImage, outDims);
            showForegroundImageView.setImageBitmap(foregroundBitmap);

        } catch (Exception e) {
            Log.e("Log", "fail to preform segmentation"); // 打印日志
            e.printStackTrace();
        }
    }


    public static String assetFilePath(Context context, String assetName) throws IOException {
        File file = new File(context.getFilesDir(), assetName); // 创建File对象
        if (file.exists() && file.length() > 0) { // 如果文件存在且文件长度大于0，则返回文件的绝对路径
            return file.getAbsolutePath();
        }

        try (InputStream is = context.getAssets().open(assetName)) {
            try (OutputStream os = new FileOutputStream(file)) { // 创建FileOutputStream对象
                byte[] buffer = new byte[4 * 1024]; // 创建缓冲区
                int read; // 创建read变量
                while ((read = is.read(buffer)) != -1) { // 读取文件
                    os.write(buffer, 0, read);
                }
                os.flush();
            }
            return file.getAbsolutePath();
        }
    }
}
